import sys
from get_Eb2c import *
from get_cylinder_xyz import *
from scipy.spatial.transform import Rotation as R
import os
def getE_l2b(source_points,target_points):
    centroid_source = np.mean(source_points, axis=0)
    centroid_target = np.mean(target_points, axis=0)

    source_points_centered = source_points - centroid_source
    target_points_centered = target_points - centroid_target

    H = np.dot(source_points_centered.T, target_points_centered)

    U, s, Vt = np.linalg.svd(H)
    R_matrix = np.dot(Vt.T, U.T)

    if np.linalg.det(R_matrix) < 0:
        Vt[:, -1] *= -1
        R_matrix = np.dot(Vt.T, U.T)

    t = -np.dot(R_matrix, centroid_source) + centroid_target

    rotation = R.from_matrix(R_matrix)
    # pdb.set_trace()
    P = np.concatenate([rotation.as_matrix(),t.reshape(3,1)],axis=1)
    P = np.vstack([P,[[0,0,0,1]]])
    return P

def mkdir_or_exist(dir_name, mode=0o777):
    if osp.exists(dir_name):
        return
    else:
        if dir_name == '':
            return
        dir_name = osp.expanduser(dir_name)
        os.makedirs(dir_name, mode=mode)

if __name__=="__main__":
    i_p=sys.argv[1]
    lidar_p=sys.argv[2]
    qtruck_name=sys.argv[3]
    cam_type=sys.argv[4]
    auto=sys.argv[5]
    config_dir = sys.argv[6] #"/home/westwell/welldriver"
    # config_dir = "/home/westwell/welldriver"

    debug_str="debug=1"
    log_dir="./log/"+qtruck_name+"/"+cam_type
    mkdir_or_exist(log_dir)

    if debug_str == "debug=1":
        debug=True
    else:
        debug=False
    # 点云 截取区域， 车头x 值为 4.85
    # xyz_min=[6, -2.1, 0.2]
    # xyz_max=[15,2.1, 1]
    if "rear_left" in cam_type:
        xyz_min=[-8,1.6, 0.2]
        xyz_max=[8, 10, 1]
    elif "rear_right" in cam_type:
        xyz_min=[-8, -10, 0.2]
        xyz_max=[8,  -1.6, 1]
    elif "front_left" in cam_type:
        xyz_min=[2,1.6, 0.2]
        xyz_max=[15, 10, 1]
    elif "front_right" in cam_type:
        xyz_min=[0, -10, 0.2]
        xyz_max=[10,  -1.6, 1]
    else:# 正前 立柱区域
        xyz_min=[14.8+2, -3, 0.2]
        xyz_max=[4.8+30, 3, 1]

    # config_dir = "/opt/qomolo/wp/"

    K_dir =config_dir+"/wellpilot_config/calib/cameras_in_calib_file"
    E_root=config_dir+"/wellpilot_config/calib/cameras_ex_calib_file"

    # code 
    # 1 get Eb2c by aruco board
    ID_min,ID_max=187,187+7
    ids_list=[]
    for i in range(ID_min,ID_max,1):
        ids_list.append(i)
        # print("ids: \n",ids_list)
    print(len(ids_list))
    dictionary = cv2.aruco.getPredefinedDictionary(dict=cv2.aruco.DICT_4X4_250)
    board = cv2.aruco.CharucoBoard((5,3),
                                squareLength=1.,
                                markerLength=0.750,
                                dictionary=dictionary,
                                ids=np.array(ids_list)
    )
    
    # ID_min,ID_max=187,187+37
    # ids_list=[]
    # for i in range(ID_min,ID_max,1):
    #     ids_list.append(i)
    #     print("ids: \n",ids_list)
    # print(len(ids_list))
    # dictionary = cv2.aruco.getPredefinedDictionary(dict=cv2.aruco.DICT_4X4_250)
    # board = cv2.aruco.CharucoBoard((15,5),
    #                             squareLength=1.,
    #                             markerLength=0.750,
    #                             dictionary=dictionary,
    #                             ids=np.array(ids_list)
    # )
    
    
    K,D=getKD(osp.join(K_dir,qtruck_name,cam_type+".yaml"))
    image = cv2.imread(i_p)
    Eb2c, img_points, obj_points,ids,img_out = get_Eb2c(K,D,image,board,dictionary,debug)
    print("get_Eb2c : \n",Eb2c)
    # np.savetxt("Eb2c.txt",Eb2c)
    cv2.imwrite(log_dir+"/img_corners_"+qtruck_name+"_"+cam_type+".jpg",img_out)

    # 2 get El2b
    # 2.1 get 4cylinder xyz
    if auto=="auto":
        auto = True
    else:
        auto = False
    pts_num = 4
    if auto:
        p_lidar = get_lidar_pts_auto(lidar_p,pts_num,xyz_min,xyz_max,cam_type,debug)
    else:
        p_lidar = get_lidar_pts_manual(lidar_p,pts_num)
    # np.savetxt("lidar_pts.txt",p_lidar)
    
    # 2.2 get board 4cylinder xyz
    # sort by xy
    '''
    *   *
    p4  p3
    *   *
    p1  p2
    '''
    p1 =[0.5,2.5,0]
    p2 =[4.5,2.5,0]
    p3 =[4.5,0.5,0]
    p4 =[0.5,0.5,0]
    # new board
    # p1 =[4.1,4.1,0]
    # p2 =[10.1,4.1,0]
    # p3 =[10.1,0.1,0]
    # p4 =[4.1,0.1,0]

    p_board = np.array([p1,p2,p3,p4])

    # 2.3
    El2b = getE_l2b(p_lidar,p_board)
    # np.savetxt("El2b.txt",El2b)

    # 3 get final El2c
    El2c= np.dot(Eb2c,El2b)
    print("El2c, \n",El2c)

    E_p=osp.join(E_root,qtruck_name,cam_type+".txt")
    mkdir_or_exist(osp.dirname(E_p))
    np.savetxt(E_p,El2c,fmt='%.8f')

    print("calib E log: ",log_dir+"/img_corners_"+qtruck_name+"_"+cam_type+".jpg")
    print("calib E save: \n",E_p)